#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2020/4/3
# @Author  : geekhch
# @Email   : geekhch@qq.com
# @Desc    : None

import sys, json
import torch
import os
import numpy as np
from OpenNRE import opennre
from OpenNRE.opennre import encoder, model, framework
from myPath import *
from myVocab import vocab
from utils import logger

# Some basic settings
if not os.path.exists(DULE_MODEL):
    os.makedirs(DULE_MODEL)

from datetime import datetime
current_time = datetime.now().strftime('%b%d_%H-%M-%S')
ckpt = path.join(DULE_MODEL, f'{current_time}_pcnn.pth.tar')

rel2id = json.load(open(os.path.join(DULE_DIR, 'rel2id.json')))
wordi2d = vocab.get_word2id_dict()
word2vec = vocab.get_embedding()

# Define the sentence encoder
# Define the sentence encoder
sentence_encoder = opennre.encoder.PCNNEncoder(
    token2id=wordi2d,
    max_length=60,
    word_size=word2vec.shape[1],
    position_size=5,
    hidden_size=300,
    blank_padding=True,
    kernel_size=3,
    padding_size=1,
    word2vec=word2vec,
    dropout=0.5
)

# Define the model
model = opennre.model.BagAttention(sentence_encoder, len(rel2id), rel2id)
logger.info(str(model))

# Define the whole training framework
framework = opennre.framework.BagRE(
    train_path=os.path.join(DULE_DIR, 'train.txt'),
    val_path=os.path.join(DULE_DIR, 'dev.txt'),
    test_path=os.path.join(DULE_DIR, 'dev.txt'),
    model=model,
    ckpt=ckpt,
    batch_size=64,
    max_epoch=60,
    lr=0.1,
    weight_decay=1e-5,
    opt='adam',)


if __name__ == '__main__':
    # Train the model
    # framework.train_model()
    print(model)
    # Test the model
